import abc
from torch.utils.data import Dataset, DataLoader
import numpy as np
import torch
import json


# 定义dataloader导入数据
def default_loader(img_path):
    img_numpy = np.load(img_path)
    img_tensor = img_numpy / 255
    img_tensor = torch.from_numpy(img_tensor).float().cuda()
    return img_tensor


# 定义获取本次实验所用的DRR的数据库
class Trainset(Dataset):
    def __init__(self, file_train, number_train, loader=default_loader):
        # 定义好 image 的路径
        self.images = file_train
        self.target = number_train
        self.loader = loader

    def __getitem__(self, index):
        fn = self.images[index]
        img = self.loader(fn)
        # 获取真实角度
        angle_x = self.target[index]["angle_x"]
        angle_y = self.target[index]["angle_y"]
        angle_z = self.target[index]["angle_z"]
        tx = self.target[index]["tx"]
        ty = self.target[index]["ty"]
        tz = self.target[index]["tz"]
        # 获取初步预测角度
        pre_rx = self.target[index]["pre_rx"]
        pre_ry = self.target[index]["pre_ry"]
        pre_rz = self.target[index]["pre_rz"]
        pre_tx = self.target[index]["pre_tx"]
        pre_ty = self.target[index]["pre_ty"]
        pre_tz = self.target[index]["pre_tz"]
        # 获取角度差
        margin_rx = angle_x - pre_rx
        margin_ry = angle_y - pre_ry
        margin_rz = angle_z - pre_rz
        margin_tx = (tx - pre_tx)/2
        margin_ty = (ty - pre_ty)/2
        margin_tz = (tz - pre_tz)/5
        target = np.array([margin_rx, margin_ry, margin_rz, margin_tx, margin_ty, margin_tz])
        target = torch.from_numpy(target).float().cuda()
        # 获取用于计算最大顶点距离的真实位姿
        tru_rx = self.target[index]["angle_x"] + 15
        tru_ry = self.target[index]["angle_y"] - 255
        tru_rz = self.target[index]["angle_z"] + 15
        tru_tx = (self.target[index]["tx"] + 30) / 2
        tru_ty = (self.target[index]["ty"] + 30) / 2
        tru_tz = (self.target[index]["tz"] + 100) / 5
        tru_pos = np.array([tru_rx, tru_ry, tru_rz, tru_tx, tru_ty, tru_tz])
        tru_pos = torch.from_numpy(tru_pos).float().cuda()
        return img, target, tru_pos

    def __len__(self):
        return len(self.images)


# 通过json获取图片的信息
def get_info_from_json(DRR_path, fn_path):
    """

    :param DRR_path: 输入与DRR图片相关的文件路径
    :param fn_path: 输入与DRR图片信息相关的文件路径
    :return:列表形式的图片路径和信息
    """
    # 载入图片数据
    img_path = []
    img_index = []
    # 读取json文件
    with open(fn_path, 'r') as img_json:
        img_info = json.load(img_json)
        # print(jsoimgn.dumps(train_info, indent=4))
    for key in img_info:
        # print(train_info[key])
        path = DRR_path + "/" + img_info[key]['name']
        img_path.append(path)
        img_index.append(img_info[key])
    return img_path, img_index


# 实例化某一个数据库，包括图片数据的导入和将数据库封装为一个迭代器
class DRR_dataset:
    def __init__(self, DRR_path, fn_path, Batch_size):
        train_path, train_index = get_info_from_json(DRR_path=DRR_path, fn_path=fn_path)
        self.train_data = Trainset(train_path, train_index)
        train_size = int(len(self.train_data) * 0.7)
        test_size = len(self.train_data) - train_size
        self.train_dataset, self.test_dataset = torch.utils.data.random_split(self.train_data, [train_size, test_size])
        self.Batch_size = Batch_size
        self.train_loader = DataLoader(self.train_dataset, batch_size=self.Batch_size, shuffle=True, drop_last=True)
        self.train_iter = iter(self.train_loader)
        self.test_loader = DataLoader(self.test_dataset, batch_size=self.Batch_size, shuffle=True, drop_last=True)
        self.test_iter = iter(self.train_loader)

    def get_train_data(self):
        try:
            batch_x, batch_y, tru_pos = next(self.train_iter)
        except StopIteration:
            self.train_loader = DataLoader(self.train_dataset, batch_size=self.Batch_size, shuffle=True, drop_last=True)
            self.train_iter = iter(self.train_loader)
            batch_x, batch_y, tru_pos = next(self.train_iter)

        return batch_x, batch_y, tru_pos

    def get_test_data(self):
        try:
            batch_x, batch_y, tru_pos = next(self.test_iter)
        except StopIteration:
            self.test_loader = DataLoader(self.test_dataset, batch_size=self.Batch_size, shuffle=True, drop_last=True)
            self.test_iter = iter(self.test_loader)
            batch_x, batch_y, tru_pos = next(self.test_iter)

        return batch_x, batch_y, tru_pos
